package org.obp.core.pool;

import com.alibaba.ttl.TtlCallable;
import com.alibaba.ttl.TtlRunnable;
import com.alibaba.ttl.threadpool.TtlExecutors;
import lombok.extern.slf4j.Slf4j;
import org.obp.core.context.ObpContext;
import org.obp.core.exception.BizException;
import org.obp.core.result.ApiResult;
import org.slf4j.MDC;

import java.util.concurrent.*;

@Slf4j
public abstract class TtlPoolManager {

    ExecutorService executorService = initExecutorService();

    ExecutorService ttlExecutorService = TtlExecutors.getTtlExecutorService(executorService);

    protected abstract TtlThreadPool initExecutorService();

    public static class TtlThreadPool extends ThreadPoolExecutor {

        public TtlThreadPool(int corePoolSize,
                             int maximumPoolSize,
                             long keepAliveTime,
                             TimeUnit unit,
                             BlockingQueue<Runnable> workQueue) {
            super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
        }

        @Override
        protected void afterExecute(Runnable r, Throwable t) {
            FutureTask<?> futureTask = (FutureTask<?>) r;
            try {
                futureTask.get();
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            } catch (ExecutionException e) {

                if (e.getCause() instanceof BizException) {
                    BizException bizException = (BizException) e.getCause();
                    MDC.put(ApiResult.TRACE_ID, bizException.getTraceId());
                    log.error(bizException.getMessage(), bizException);

                } else if (e.getCause() instanceof Exception) {
                    Exception exception = (Exception) e.getCause();
                    log.error(exception.getMessage(), exception);
                }
            } finally {
                MDC.remove(ApiResult.TRACE_ID);
            }
            super.afterExecute(r, t);
        }

    }

    public Future<?> submit(Runnable task) {
        return ttlExecutorService.submit(() -> {
            try {
                MDC.put(ApiResult.TRACE_ID, getTraceId());
                TtlRunnable.get(task).run();
            } catch (Exception e) {
                log.error(e.getMessage(), e);
            } finally {
                MDC.remove(ApiResult.TRACE_ID);
            }
        });
    }

    public <T> Future<T> submit(Callable<T> task) {
        return ttlExecutorService.submit(() -> {
            try {
                MDC.put(ApiResult.TRACE_ID, getTraceId());
                return TtlCallable.get(task).call();
            } finally {
                MDC.remove(ApiResult.TRACE_ID);
            }
        });
    }

    public static String getTraceId() {
        return ObpContext.getTraceId();
    }

    public ExecutorService getTtlExecutorService() {
        return ttlExecutorService;
    }

}
